import os
import numpy as np
from utils.paths import is_contain_file


def load_daily_sleep30(fp_sleep30):
    lines = []
    with open(fp_sleep30, 'r') as h:
        lines = [line.rstrip('\r\t\n ') for line in h.readlines()]
    sleep30 = []
    for line in lines[1:]:
        if len(line) < 1:
            continue
        _, val = [float(v) for v in line.split('\t')]
        sleep30.append(val)
    return sleep30


def save_stat_sleep(fp_out, chn_list, sleep30_stack):
    '''
        将sleep30数据汇总保存到csv文件中
        param.chn_list: n_mouse 通道号名称
        param.sleep30_stack: n_mouse x 48 小鼠30min睡眠24小时数据顺序与chn_list对应
    '''
    # 确认封装为mat
    if isinstance(sleep30_stack, list):
        sleep30_stack = np.array(sleep30_stack)
    # 计算白天、黑夜、总量睡眠
    day_sleeps, night_sleeps = np.sum(sleep30_stack[:, :24], axis=1), np.sum(sleep30_stack[:, 24:], axis=1)
    total_sleeps = day_sleeps + night_sleeps
    with open(fp_out, 'w') as h:
        # 输出表头
        h.write('channel\tday_sleep\tnight_sleep\ttotal_sleep (unit:sec)\n')
        # 输出每只鼠的睡眠数据
        zipped_data = zip(chn_list, day_sleeps, night_sleeps, total_sleeps)
        for chn, day_sleep, night_sleep, total_sleep in zipped_data:
            h.write(chn) # 输出对应通道号
            [h.write('\t%.4f' % v) for v in [day_sleep, night_sleep, total_sleep]] # 连续白天、黑夜、总量睡眠
            h.write('\n')

def summarise_section_sleep_stat(root_path, f_daily_sleep):
    chn_list, sleep30_stack = [], []
    for chn in os.listdir(root_path):
        dpath = os.path.join(root_path, chn) # 获取子目录绝对路径
        if os.path.isdir(dpath) and is_contain_file(dpath, f_daily_sleep):
            fp_daily_sleep = os.path.join(dpath, f_daily_sleep) # 获取目标文件绝对路径
            sleep30 = load_daily_sleep30(fp_daily_sleep) # 读取24小时48点数据
            chn_list.append(chn) # 记录当前通道名称
            sleep30_stack.append(sleep30) # 记录当前数据
    # 输出汇总的数据
    fp_out = os.path.join(root_path, 'summary_of_' + f_daily_sleep)
    save_stat_sleep(fp_out, chn_list, sleep30_stack)


if __name__ == '__main__':
    root_path = '/home/youngway/data/mtrace/20201202_c57-males-dvr2'
    f_daily_sleep_list =  [prefix + '_daily_sleep.txt' for prefix in ['before_inject', 'after_1stinject', 'after_2ndinject']]
    for f_daily_sleep in f_daily_sleep_list:
        summarise_section_sleep_stat(root_path, f_daily_sleep)